package cn.com.config;

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.sql.DataSource;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;


public class UserRoutingDataSource extends AbstractRoutingDataSource {


    private final Map<Object, Object> userDataSources = new ConcurrentHashMap<>();
    private final Map<Object, DataSource> resolvedUserDataSources = new ConcurrentHashMap<>();

    public void registerUserDataSource(Object id, Object dataSource) {
        userDataSources.put(id, dataSource);
        resolveDataSource(id, dataSource);
    }

    /**
     * 重写resolveDataSource，动态解析数据源
     * @param id
     * @param dataSource
     */
    public void resolveDataSource(Object id, Object dataSource) {
        Object lookupKey = this.resolveSpecifiedLookupKey(id);
        DataSource resolveDataSource = this.resolveSpecifiedDataSource(dataSource);
        resolvedUserDataSources.put(lookupKey, resolveDataSource);
    }

    public Object lookUpUserDataSource(Object id) {
        return userDataSources.get(id);
    }

    public void removeDataSource(Object id) {
        userDataSources.remove(id);
    }

    @Override
    public Map<Object, DataSource> getResolvedDataSources() {
        return resolvedUserDataSources;
    }

    @Override
    public void afterPropertiesSet() {
        this.setTargetDataSources(userDataSources);
        userDataSources.forEach(this::resolveDataSource);
    }

    @Override
    protected Object determineCurrentLookupKey() {
        return UserDataSourceContext.getUserId();
    }

    /**
     * 重写determineTargetDataSource，动态选择目标数据源
     * @return
     */
    @Override
    protected DataSource determineTargetDataSource() {
        Object userId = this.determineCurrentLookupKey();
        DataSource userDataSource = resolvedUserDataSources.get(userId);
        if(Objects.nonNull(userDataSource)) {
            return userDataSource;
        }
        throw new RuntimeException("该用户没有初始化数据源");
    }
}
